from typing import Optional

from IPython.core.display import Image
from IPython.core.display_functions import display
from langgraph.graph import StateGraph
from langgraph.graph import START, END
from pydantic import BaseModel


class LoopState(BaseModel):
    x: int


def increment(state: LoopState) -> LoopState:
    print(f"[increment]Received x=: {state.x}")
    return LoopState(x=state.x + 1)


def is_done(state: LoopState) -> bool:
    return state.x > 10


if __name__ == '__main__':
    # pip install pyppeteer ipython
    builder = StateGraph(LoopState)
    builder.add_node("increment", increment)
    # 添加循环控制。
    builder.add_conditional_edges("increment", is_done, {
        True: END,
        False: "increment",
    })
    builder.add_edge(START, "increment")
    print(builder.edges)
    print(builder.nodes)
    graph = builder.compile()
    graph_image = graph.get_graph(xray=True).draw_mermaid_png()
    with open("graph_loop.png", "wb") as f:
        f.write(graph_image)
    print("Graph saved to graph.png")
    print(f"当前值: {graph.invoke(LoopState(x=5))}")
